/*********************
* author:yangxin
* date :2025.3.10
* test wht c++ and avx2
**************************/
#include <iostream>
#include <chrono>
#include <cstring>
#include <intrin.h>
#include <malloc.h>
#include <immintrin.h>
#include <assert.h>


using namespace std;

//declare asm wht
extern "C" void fwht_8x8_row_pass_avx2(const short* src, short* dst, int row, int dpit);

//detect support avx2
bool is_avx2_supported(){
    int cpuInfo[4];

    __cpuid(cpuInfo, 1);
    bool avx_supported = (cpuInfo[2]&(1<<28)) != 0;
    cout << "avx: " << avx_supported << endl;
    if(!avx_supported)
        return false;

    __cpuid(cpuInfo, 7);
    bool avx2_supported = (cpuInfo[1]&(1<<5)) != 0;
    cout << "avx2: " << avx2_supported << endl;

    return avx2_supported;
}


//c++ wht
void fwht_8x8_row_pass(const short* src, short* dst, int row, int dpit) {
    short* sptr = (short*)src;
    int j;

    for(j = 0; j < row; j ++ ){
        int adata[8];
        int bdata[8];
        bdata[0] = sptr[0] + sptr[4];
        bdata[1] = sptr[1] + sptr[5];
        bdata[2] = sptr[2] + sptr[6];
        bdata[3] = sptr[3] + sptr[7];
        bdata[4] = sptr[0] - sptr[4];
        bdata[5] = sptr[1] - sptr[5];
        bdata[6] = sptr[2] - sptr[6];
        bdata[7] = sptr[3] - sptr[7];

        adata[0] = bdata[0] + bdata[2];
        adata[1] = bdata[1] + bdata[3];
        adata[2] = bdata[0] - bdata[2];
        adata[3] = bdata[1] - bdata[3];
        adata[4] = bdata[4] + bdata[6];
        adata[5] = bdata[5] + bdata[7];
        adata[6] = bdata[4] - bdata[6];
        adata[7] = bdata[5] - bdata[7];

        dst[0*dpit] = adata[0] + adata[1];
        dst[1*dpit] = adata[0] - adata[1];
        dst[2*dpit] = adata[2] + adata[3];
        dst[3*dpit] = adata[2] - adata[3];
        dst[4*dpit] = adata[4] + adata[5];
        dst[5*dpit] = adata[4] - adata[5];
        dst[6*dpit] = adata[6] + adata[7];
        dst[7*dpit] = adata[6] - adata[7];


        dst++; 
        sptr += 8;
    }
}



typedef void (*fwht_Function)(const short*, short*, int, int);

// 打印8x8矩阵
void print_matrix(const short* data, int dpit) {
    for(int i = 0; i < 8; i++) {
        for(int j = 0; j < 8; j++) {
            cout << data[j*dpit + i] << "\t";
        }
        cout << "\n";
    }
}

//test
void testFunction() {

    const int row = 8;
    const int col = 8;
    const int dpit = 8;

    // 初始化测试数据
    alignas(32) short src[64] = {0};
    for(int i = 0; i < 8; ++i) {
        src[i] = i;  
    }
    cout << "Input:\n";
    print_matrix(src, 8);

    // c++
    alignas(32) short dst_orig[64] = {0};
    fwht_8x8_row_pass(src, dst_orig, row, dpit);
    cout << "\nOriginal c++ Version Result:\n";
    print_matrix(dst_orig, 8);

    // avx2
    if(is_avx2_supported()){

        alignas(32) short dst_avx2[64] = {0};
        fwht_Function fwht_avx2 = fwht_8x8_row_pass_avx2;
        fwht_avx2(src, dst_avx2, 8, 8);

        cout << "\nAVX2 Version Result:\n";
        print_matrix(dst_avx2, 8);

    }else{
        cout << " is not support avx2!" << endl;
    }


}



int main() {
    try{
        testFunction();
        
    }catch (const std::exception& e) {
        std::cerr << "Exception caught: " << e.what() << std::endl;
        return 1;
    }


    return 0;
}